from enum import Enum
from typing import Optional, Dict
from service.chat.struct_info import InputsInfo, CompleteInfo, DetailMathInfo, SolveInputs, SolverResults, MathKnowledgeResult


class ChatState(Enum):
    INIT = 1
    EXTRACT_INFO = 2
    DETAIL_INFO = 3
    SOLVE_PROBLEM = 4
    EXPLAIN_KNOWLEDGE = 5
    COMPLETE = 6


class ChatStatus:

    def __init__(self):
        self.reset()

    def reset(self):
        self.chat_state = ChatState.INIT
        self.inputs_info = None
        self.detail_info = None
        self.problem_result = None
        self.knowledge_result = None

    def get_current_state(self):
        return self.chat_state
    
    def get_current_inputs(self):
        if self.chat_state == ChatState.INIT:
            return None
        elif self.chat_state == ChatState.EXTRACT_INFO:
            return self.inputs_info
        elif self.chat_state == ChatState.DETAIL_INFO:
            return SolveInputs(self.inputs_info.latex, self.detail_info)
        elif self.chat_state == ChatState.SOLVE_PROBLEM:
            return self.problem_result
        elif self.chat_state == ChatState.EXPLAIN_KNOWLEDGE:
            return self.knowledge_result
        elif self.chat_state == ChatState.COMPLETE:
            return self.complete_info
        else:
            return None

    def update_current_state(self, new_state: ChatState, value_info) -> None:
        self.chat_state = new_state

        if self.chat_state == ChatState.INIT:
            pass
        elif self.chat_state == ChatState.EXTRACT_INFO:
            self.inputs_info : InputsInfo = value_info
        elif self.chat_state == ChatState.DETAIL_INFO:
            self.detail_info : DetailMathInfo = value_info
        elif self.chat_state == ChatState.SOLVE_PROBLEM:
            self.problem_result : SolverResults = value_info
        elif self.chat_state == ChatState.EXPLAIN_KNOWLEDGE:
            self.knowledge_result : MathKnowledgeResult = value_info
        elif self.chat_state == ChatState.COMPLETE:
            self.complete_info : CompleteInfo = value_info
        else:
            pass

class ChatManager:

    def __init__(self):
        self.chatmessages: Dict[str, ChatStatus] = {}

    def get_current_inputs(self, session_id: str):
        if session_id in self.chatmessages:
            return self.chatmessages[session_id].get_current_inputs()
        else:
            self.chatmessages[session_id] = ChatStatus()
            return self.chatmessages[session_id].get_current_inputs()

    def get_current_state(self, session_id: str)-> ChatState:
        if session_id in self.chatmessages:
            return self.chatmessages[session_id].get_current_state()
        else:
            self.chatmessages[session_id] = ChatStatus()
            return self.chatmessages[session_id].get_current_state()

    def update_current_state(self, session_id: str, new_state: ChatState, value_info) -> None:
        if session_id in self.chatmessages:
            self.chatmessages[session_id].update_current_state(new_state, value_info)
        else:
            self.chatmessages[session_id] = ChatStatus()
            return self.chatmessages[session_id].update_current_state(new_state, value_info)
        
    def reset(self, session_id: str)-> None:
        if session_id in self.chatmessages:
            self.chatmessages[session_id].reset()
            del self.chatmessages[session_id]
        # else:
        #     self.chatmessages[session_id] = ChatStatus()
        #     return self.chatmessages[session_id].reset()

    @classmethod
    def get_instance(cls):
        if not hasattr(cls, 'instance'):
            cls.instance = cls()
        return cls.instance
